package test.demo;

import cn.hutool.core.util.ArrayUtil;
import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;
import com.alibaba.fastjson.JSON;

import java.util.Date;
import java.util.HashMap;
import java.util.Map;

/**
 * Green & Ampt
 *
 * @author 初京刚
 * @version 1.0
 */

public class GreenAmpt extends PredictBase {

    /**
     * 当前土壤类型、土壤覆被下水分条件II（一般湿润）时的SCS曲线数（Curve Number）
     */
    private double CN2;

    /**
     * 前期土壤水分条件（1-干燥，2-一般湿润，3-湿润）
     */
    private int AMC;

    /**
     * 流域平均坡度，%
     */
    private double slp;

    /**
     * 降水量
     */
    private double[] floodDrp;

    /**
     * 降水时间
     */
    private Date[] floodTm;

    /**
     * 预报开始时间
     */
    private Date startTm;

    /**
     * 预报开始数组下标
     */
    private int startIndex;

    /**
     * 降水序列长度
     */
    private Integer floodRainRange;

    /**
     * 降水时段长
     */
    private double timeInterval;

    /**
     * 时段长，minutes
     */
    private double idt;

    /**
     * 最大时段数
     */
    private int mstep;

    /**
     * 初始时段下渗率，mm/hr
     */
    private double newrti;

    /**
     * 时段数
     */
    private int nstep;

    /**
     * 土壤饱和传导率，mm/hr
     */
    private double sol_k;

    /**
     * 土壤孔隙度
     */
    private double sol_por;

    /**
     * 土壤田间持水量，mm
     */
    private double sol_sumfc;

    /**
     * 初始土壤蓄水量，mm
     */
    private double sol_sw;

    /**
     * 降雨事件标志，0-午夜无降雨事件，1-午夜有降雨事件
     */
    private int swtrg;

    /**
     * 湿润前沿的平均毛细管吸力，mm
     */
    private double wfsh;

    /**
     *小时累积地表径流，mm
     */
    private double[] hhqday;

    /**
     * 地表径流，mm
     */
    private double surfq;

    /**
     * 时段累积下渗，mm
     */
    private double[] cuminf;

    /**
     * 时段累积降雨，mm
     */
    private double[] cumr;

    /**
     * 时段累积径流，mm
     */
    private double[] excum;

    /**
     * 时段径流，mm
     */
    private double[] exinc;

    /**
     * 时段下渗率，mm/hr
     */
    private double[] rateinf;

    /**
     * 时段降雨强度，mm/hr
     */
    private double[] rintns;

    public GreenAmpt(Map<String, String> data) {
        super();
        CN2 = Double.parseDouble(data.get("CN2"));
        slp = Double.parseDouble(data.get("SLOPE"));
        AMC = Integer.parseInt(data.get("AMC"));
        floodTm = JSON.parseObject(data.get("dt"), Date[].class);               // input[i].data.dt
        startTm = JSON.parseObject(data.get("start"), Date.class);
        startIndex = ArrayUtil.indexOf(floodTm, startTm);
        floodTm = JSON.parseObject(data.get("dt"), Date[].class);   // input[i].data.dt
        floodDrp = JSON.parseObject(data.get("rain"), double[].class);// input[i].data.rain
        floodRainRange = floodDrp.length;
        if (floodTm.length > 1) {
            timeInterval = getHours(floodTm[0], floodTm[1]);
        } else {
            timeInterval = Double.parseDouble(data.get("clen"));                  // index.clen
        }
        idt = timeInterval * 60;
        mstep = floodRainRange;
        nstep = floodRainRange;
        newrti = Double.parseDouble(data.get("RTI"));;
        sol_k = Double.parseDouble(data.get("K"));
        sol_por = Double.parseDouble(data.get("POR"));
        sol_sumfc = Double.parseDouble(data.get("SUMFC"));
        sol_sw = Double.parseDouble(data.get("SW"));
        swtrg = Integer.parseInt(data.get("TRG"));
        wfsh = Double.parseDouble(data.get("WFSH"));
        hhqday = new double[floodRainRange * (int)timeInterval];
        cuminf = new double[floodRainRange + 1];
        cumr = new double[floodRainRange + 1];
        excum = new double[floodRainRange + 1];
        exinc = new double[floodRainRange + 1];
        rateinf = new double[floodRainRange + 1];
        rintns = new double[floodRainRange + 1];
    }

    /**
     * 模型计算
     */
    public Map<String, double[]> SCSmodel() {

        // Green & Ampt 产流
        GA();

        Map<String, double[]> simulationResult = new HashMap<>();
        simulationResult.put("runoffSim", exinc);
        simulationResult.put("rain", floodDrp);
        return simulationResult;
    }

    /**
     * Green & Ampt 产流模型
     */
    private void GA() {
        double CN2S, CN1, CN3, CN;
        CN1 = CN2 - 20 * (100 - CN2) / (100 - CN2 + Math.exp(2.533 - 0.0636 * (100 - CN2)));
        CN3 = CN2 * Math.exp(0.00673 * (100 - CN2));
        CN2S = (CN3 - CN2) / 3 * (1 - 2 * Math.exp(-13.86 * slp)) + CN2;
        CN = CN2S;
        if(AMC == 1) {
            CN = CN2S - 20 * (100 - CN2S) / (100 - CN2S + Math.exp(2.533 - 0.0636 * (100 - CN2S)));
        } else if(AMC == 3) {
            CN = CN2S * Math.exp(0.00673 * (100 - CN2S));
        }
        if(sum(floodDrp) > 0.1){
            double adj_hc = 0.0;
            adj_hc = (56.82 * Math.pow(sol_k, 0.286)) / (1 + 0.051 * Math.exp(0.062 * CN)) - 2;
            if(adj_hc <= 0.0) {
                adj_hc = 0.001;
            }
            double dthet = 0.0;
            if(swtrg == 1) {
                swtrg = 0;
                dthet = 0.001 * sol_por * 0.95;
                rateinf[0] = newrti;
                newrti = 0.0;
            } else {
                double soilw =0.0;
                if(sol_sw >= sol_sumfc) {
                    soilw = 0.999 * sol_sumfc;
                } else {
                    soilw = sol_sw;
                }
                dthet = (1 - soilw / sol_sumfc) * sol_por * 0.95;
                rateinf[0] = 2000.0;
            }
            double psidt = 0.0;
            psidt = dthet * wfsh;
            int k = 0;
            rintns[0] = floodDrp[1] / idt * 60;
            cumr[0] = 0.0;
            cuminf[0] = 0.0;
            excum[0] = 0.0;
            exinc[0] = 0.0;
            for(k = 1; k < nstep + 1; k++) {
                cumr[k] = cumr[k - 1] + floodDrp[k];
                rintns[k] = floodDrp[k + 1] / idt * 60;
                if(rateinf[k - 1] >= rintns[k - 1]) {
                    cuminf[k] = cuminf[k - 1] + rintns[k - 1] * idt / 60;
                    if(excum[k - 1] > 0.0) {
                        excum[k] = excum[k - 1];
                    } else {
                        excum[k] = 0.0;
                    }
                    exinc[k] = 0.0;
                } else {
                    double tst = 0.0;
                    tst = adj_hc * idt / 60;
                    double f1 = 0.0;
                    f1 = cuminf[k - 1] + adj_hc * idt / 60 + psidt * Math.log((tst + psidt) / (cuminf[k - 1] + psidt));
                    while(Math.abs(f1 - tst) > 0.001) {
                        tst = 0.0;
                        tst = f1;
                        f1 = cuminf[k - 1] + adj_hc * idt / 60 + psidt * Math.log((tst + psidt) / (cuminf[k - 1] + psidt));
                    }
                    cuminf[k] = f1;
                    excum[k] = cumr[k] - cuminf[k];
                    exinc[k] = excum[k] - excum[k - 1];
                    if(exinc[0] < 0) {
                        exinc[0] = 0.0;
                    }
                    int kk =0;
                    kk = (k - 1) * (int)idt;
                    if(Math.floorMod(kk, 60) == 0) {
                        kk = kk / 60;
                    } else {
                        kk = 1 + kk / 60;
                    }
                    hhqday[kk] += exinc[k];
                    surfq += excum[k];
                }
                rateinf[k] = adj_hc * (1 + psidt / (cuminf[k] + 1.0E-6));
            }
        }
    }

    private int getHours(Date dt_beg, Date dt_end) {
        int hours = -1;
        hours = Long.valueOf((dt_end.getTime() - dt_beg.getTime()) / (3600 * 1000L)).intValue();
        return hours;
    }

    private double sum(double[] data){
        double sum = 0.0;
        if(data != null){
            for(int i = 0; i < data.length; i++){
                sum += data[i];
            }
        }
        return sum;
    }

    @Override
    public void init() {

    }

    @Override
    public PredictResult predict() {
        PredictResult result = new PredictResult();
        Map<String, double[]> map = SCSmodel();

        result.setRunoffSim(map.get("runoffSim"));
        return result;
    }
}